package cli.command;

import javax.crypto.*;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Random;
import java.util.Scanner;

public class Secure implements Command {
    private Scanner in;
    private PrintWriter out;
    private Path currentPath;
    private boolean e = false; // 是否加密
    private boolean d = false; // 是否解密
    private boolean f = false; // 是否指定dst文件路径
    private String src = null;
    private String dst = null;

    public Secure(Scanner in, PrintWriter out, Path currentPath) {
        this.in = in;
        this.out = out;
        this.currentPath = currentPath;
    }

    @Override
    public void handle(String[] args) throws Exception {
        accept(args);
        execute();
    }

    private void accept(String[] args) {
        int pos = 0;
        if (pos < args.length) {
            if (args[pos].startsWith("-")) {
                e = args[pos].contains("e");
                d = args[pos].contains("d");
                f = args[pos].contains("f");
                pos++;
            }
        }
        if (pos < args.length) {
            if (f) {
                dst = args[pos++];
            }
        }
        if (pos < args.length) {
            src = args[pos++];
        } else {
            out.println("请输入源文件路径！");
            return;
        }
    }

    private void execute() {
        Path src = currentPath.resolve(this.src);
        Path dst;
        if (this.dst == null) {
            if (e) {
                dst = currentPath.resolve(src.getFileName() + ".aes");
            } else if (d) {
                String srcName = src.getFileName().toString();
                String dstName = src.endsWith(".aes") ? srcName.substring(0, srcName.length() - 4) : srcName;
                dst = currentPath.resolve(dstName);
            } else {
                return;
            }
        } else {
            dst = currentPath.resolve(this.dst);
        }
        if (e) {
            encrypt(src, dst);
        } else if (d) {
            decrypt(src, dst);
        }
    }

    private String inputPassword() {
        out.println("请输入密码：");
        String pwd = in.nextLine();
        out.println("请确认密码：");
        String pwd2 = in.nextLine();
        if (pwd.equals(pwd2))
            return pwd;
        out.println("两次密码不一致！");
        return null;
    }

    private static Key generateKey(char[] password) throws NoSuchAlgorithmException, InvalidKeySpecException {
        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        KeySpec spec = new PBEKeySpec(password, "20215521翟毅鑫2".getBytes(StandardCharsets.UTF_8), 65536, 256);
        SecretKey secret = factory.generateSecret(spec);
        return new SecretKeySpec(secret.getEncoded(), "AES");
    }

    private void encrypt(Path src, Path dst)  {
        String pwd = inputPassword();
        if (pwd == null)
            return;
        try {
            Cipher c = Cipher.getInstance("AES");
            c.init(Cipher.ENCRYPT_MODE, generateKey(pwd.toCharArray()));
            InputStream is = Files.newInputStream(src, StandardOpenOption.READ);
            CipherInputStream cis = new CipherInputStream(is, c);
            OutputStream os = Files.newOutputStream(dst, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = cis.read(buffer)) != -1) {
                os.write(buffer, 0, bytesRead);
            }
            is.close();
            cis.close();
            os.close();
        } catch (NoSuchAlgorithmException ex) {
            ex.printStackTrace();
        } catch (NoSuchPaddingException ex) {
            ex.printStackTrace();
        } catch (InvalidKeyException ex) {
            out.println(pwd+"密码无效！");
        } catch (IOException ex) {
            out.println("加密失败！");
        } catch (InvalidKeySpecException ex) {
            ex.printStackTrace();
        }
    }

    private void decrypt(Path src, Path dst) {
        String pwd = inputPassword();
        if (pwd == null)
            return;
        try {
            Cipher c = Cipher.getInstance("AES");
            c.init(Cipher.DECRYPT_MODE, generateKey(pwd.toCharArray()));
            InputStream is = Files.newInputStream(src, StandardOpenOption.READ);
            CipherInputStream cis = new CipherInputStream(is, c);
            OutputStream os = Files.newOutputStream(dst, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = cis.read(buffer)) != -1) {
                os.write(buffer, 0, bytesRead);
            }
            is.close();
            os.close();
            cis.close();
        } catch (NoSuchAlgorithmException ex) {
            ex.printStackTrace();
        } catch (NoSuchPaddingException ex) {
            ex.printStackTrace();
        } catch (InvalidKeyException ex) {
            out.println(pwd+"密码无效！");
        } catch (IOException ex) {
            out.println("解密失败！");
        } catch (InvalidKeySpecException ex) {
            ex.printStackTrace();
        }
    }
}
